package com.shujia.mllib

import org.apache.spark.ml.classification.LogisticRegressionModel
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo05ImagePredict {
  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .appName("Demo05ImagePredict")
      .master("local[*]")
      .config("spark.sql.shuffle.partitions", 8)
      .getOrCreate()

    // 1、加载数据 并 进行特征工程处理
    val imageDF: DataFrame = spark
      .read
      .format("image")
      .load("Spark/data/mllib/data/image")

    imageDF.printSchema()
    import spark.implicits._

    val imageDataDF: DataFrame = imageDF
      .select($"image.origin", $"image.data")
      .as[(String, Array[Byte])]
      .map {
        case (filePath: String, binArr: Array[Byte]) =>
          val intArr: Array[Int] = binArr.map(_.toInt)
          val arr1or0: Array[Double] = intArr.map(int => {
            if (int < 0) {
              1.0
            } else {
              0.0
            }
          })
          val fileName: String = filePath.split("/").last
          (fileName, Vectors.dense(arr1or0).toSparse)
      }.toDF("fileName", "features")

    // 2、加载模型
    val lrModel: LogisticRegressionModel = LogisticRegressionModel.load("Spark/data/mllib/image/model")

    // 3、将经过特征工程处理好的数据带入模型 进行预测
    lrModel
      .transform(imageDataDF)
      .show()
  }

}
